#include "col_2_im_custom_tiling.h"
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_col_2_im_custom.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void col_2_im_custom(GM_ADDR col, GM_ADDR img, Col2ImCustomTilingData tiling);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;
    size_t tilingSize = 2 * sizeof(uint32_t) + 8 * sizeof(int32_t);

    int32_t N = 1;
    int32_t C = 1;
    int32_t H = 8;
    int32_t W = 256;
    int32_t kernel_H = 2;
    int32_t kernel_W = 2;
    int32_t stride_val = 1;
    int32_t padding_val = 0;

    // 计算col矩阵尺寸（输入）
    int32_t out_H = (H + 2 * padding_val - kernel_H) / stride_val + 1;
    int32_t out_W = (W + 2 * padding_val - kernel_W) / stride_val + 1;
    int32_t L = out_H * out_W;
    int32_t input_channels = C * kernel_H * kernel_W;
    size_t inputByteSize = N * input_channels * L * sizeof(float);      // col矩阵
    size_t outputByteSize = N * C * H * W * sizeof(float);              // 还原的图像

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *tiling = (uint8_t *)AscendC::GmAlloc(tilingSize);
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);
    uint8_t *col = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *img = (uint8_t *)AscendC::GmAlloc(outputByteSize);

    ReadFile("./input/input_col.bin", inputByteSize, col, inputByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);

    ICPU_RUN_KF(col_2_im_custom, blockDim, col, img,
                *reinterpret_cast<Col2ImCustomTilingData *>(tiling)); // use this macro for cpu debug

    WriteFile("./output/output_img.bin", img, outputByteSize);

    AscendC::GmFree((void *)col);
    AscendC::GmFree((void *)img);
    AscendC::GmFree((void *)tiling);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    Col2ImCustomTilingData *tiling;
    uint8_t *colHost, *imgHost;
    uint8_t *colDevice, *imgDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tiling), tilingSize));
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);

    CHECK_ACL(aclrtMallocHost((void **)(&colHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&imgHost), outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&colDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&imgDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_col.bin", inputByteSize, colHost, inputByteSize);

    CHECK_ACL(aclrtMemcpy(colDevice, inputByteSize, colHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    ACLRT_LAUNCH_KERNEL(col_2_im_custom)(blockDim, stream, colDevice, imgDevice, tiling);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(imgHost, outputByteSize, imgDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_img.bin", imgHost, outputByteSize);

    CHECK_ACL(aclrtFree(colDevice));
    CHECK_ACL(aclrtFree(imgDevice));
    CHECK_ACL(aclrtFreeHost(colHost));
    CHECK_ACL(aclrtFreeHost(imgHost));
    CHECK_ACL(aclrtFreeHost(tiling));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}